from rl_core import RL_Trainer
from task4_env import ThrottleHeightEnv  # 任务4的环境类
import os

# 定义模型保存路径和名称
model_save_path = "./models/"
model_name = "task4_ppo_model"
final_model_path = f"{model_save_path}/{model_name}_final"

# 检查模型是否已经保存
if os.path.exists(final_model_path + ".zip"):
    # 如果模型已保存，加载模型
    trainer = RL_Trainer(
        env_class=ThrottleHeightEnv,
        algorithm="PPO"
    )
    trainer.load(final_model_path)
    print(f"Loaded existing model from {final_model_path}")
else:
    # 如果模型未保存，创建新的模型
    trainer = RL_Trainer(
        env_class=ThrottleHeightEnv,
        algorithm="PPO",
        learning_rate=0.0003,  # 可根据需要调整超参数
        n_steps=2048,
        batch_size=64,
        gamma=0.95,
        gae_lambda=0.95,
        ent_coef=0.01,
        vf_coef=0.5,
        max_grad_norm=0.5
    )
    print("Created a new model.")

# 训练模型
trainer.train(
    total_timesteps=1000,  # 总训练步数
    save_path=model_save_path,  # 模型保存路径
    checkpoint_freq=100,  # 检查点保存频率
    model_name=model_name  # 模型名称
)

# 评估模型
mean_reward = trainer.evaluate(episodes=10)
print(f"平均奖励: {mean_reward}")